Add Oriented Bounding Box (OBB) support for rotated object detection by farukalamai · Pull Request #921 · roboflow/rf-detr

farukalamai · 2026-04-03T19:40:35Z

What does this PR do?

Adds OBB (Oriented Bounding Box) support to RF-DETR, enabling rotated object detection with [cx, cy, w, h, angle] box format. This includes rotated box math utilities, DOTA v1.0 dataset loader, angle prediction head, Gaussian-based matching/loss functions, and oriented postprocessing.

When oriented=True is set in ModelConfig, the model predicts a 5th dimension (rotation angle in radians) per box, uses KLD loss for regression, GWD cost for Hungarian matching, and outputs rotated corner points.

Related Issue(s): Fixes #56

Type of Change

New feature (non-breaking change that adds functionality)

Testing

I have tested this change locally
I have added/updated tests for this change

Test details:

75 new tests covering all OBB components:

test_rotated_box_ops.py (34 tests) — angle normalization, box conversions, roundtrips, GWD/KLD/ProbIoU losses, gradient flow, edge cases (zero-size boxes, large boxes)
test_dota_detection.py (21 tests) — annotation parsing, dataset loading, filtering, normalization, empty/missing files
test_obb_head.py (8 tests) — detection head output shapes, angle range, gradient flow
test_obb_matcher_criterion.py (6 tests) — oriented matching, KLD loss computation
test_obb_postprocess.py (5 tests) — output keys, shapes, scaling, batch support
test_obb_export.py (2 tests) — ONNX export produces 5D output with valid angles
test_obb_integration.py (5 tests) — config flag, namespace forwarding, dataset_file acceptance

All existing tests pass unchanged.

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my own code
My changes generate no new warnings or errors
I have updated the documentation accordingly (if applicable)

Additional Context

Key design decisions:

Separate angle_embed MLP rather than extending bbox_embed to 5D — keeps spatial box prediction unchanged
KLD loss for regression — handles angle boundary discontinuity, aspect-ratio adaptive
GWD cost for Hungarian matching — works even with zero-overlap boxes
ProbIoU for IoU-aware classification loss — replaces axis-aligned box_iou when oriented
oriented: bool = False on ModelConfig — zero impact on existing non-oriented models

Files changed

New files:

src/rfdetr/utilities/rotated_box_ops.py — box conversions, Gaussian encoding, GWD/KLD/ProbIoU
src/rfdetr/datasets/dota_detection.py — DOTA v1.0 dataset loader with annotation parser
tests/ — 7 new test files (75 tests)

Modified files:

src/rfdetr/config.py — added oriented flag and "dota" dataset option
src/rfdetr/_namespace.py — forward oriented to namespace
src/rfdetr/models/_types.py — added oriented to BuilderArgs protocol
src/rfdetr/models/heads/detection.py — added angle_embed MLP when oriented
src/rfdetr/models/lwdetr.py — angle prediction in forward pass, zero-init, builder wiring
src/rfdetr/models/matcher.py — GWD pairwise cost for oriented matching
src/rfdetr/models/criterion.py — KLD loss for oriented boxes, ProbIoU for ia_bce_loss
src/rfdetr/models/postprocess.py — oriented output with boxes_obb and corners
src/rfdetr/datasets/__init__.py — registered DOTA dataset builder

… return type

codecov · 2026-04-03T19:43:48Z

Codecov Report

❌ Patch coverage is 95.10490% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 80%. Comparing base (658f199) to head (e81c9f4).
⚠️ Report is 19 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff           @@
##           develop   #921    +/-   ##
=======================================
+ Coverage       79%    80%    +1%     
=======================================
  Files           97     99     +2     
  Lines         7793   8061   +268     
=======================================
+ Hits          6148   6410   +262     
- Misses        1645   1651     +6

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…error

…dules

Copilot

Pull request overview

Adds Oriented Bounding Box (OBB) support across the RF-DETR training/inference stack to enable rotated object detection using [cx, cy, w, h, angle] boxes.

Changes:

Introduces rotated box math utilities (corner conversion + Gaussian-based GWD/KLD/ProbIoU).
Wires an oriented flag through config/namespace/model, adding an angle prediction head and oriented matching/loss/postprocess paths.
Adds a DOTA v1.0 dataset loader and a broad new unit test suite for OBB components.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`src/rfdetr/utilities/rotated_box_ops.py`	Adds rotated box conversions and Gaussian-based similarity/loss utilities.
`src/rfdetr/datasets/dota_detection.py`	New DOTA dataset loader + transforms/normalization for OBB training.
`src/rfdetr/models/matcher.py`	Adds oriented matching path using GWD cost and `boxes_obb`.
`src/rfdetr/models/criterion.py`	Adds oriented loss path using KLD and ProbIoU for IoU-aware classification.
`src/rfdetr/models/postprocess.py`	Adds oriented postprocessing output (`boxes_obb` + `corners`).
`src/rfdetr/models/lwdetr.py`	Adds angle head + forward/export plumbing + passes `oriented` into postprocess.
`src/rfdetr/models/heads/detection.py`	Adds optional angle prediction in the standalone detection head.
`src/rfdetr/models/_types.py`	Extends builder args protocol to include `oriented`.
`src/rfdetr/datasets/__init__.py`	Registers DOTA dataset builder.
`src/rfdetr/config.py`	Adds `ModelConfig.oriented` and `TrainConfig.dataset_file="dota"`.
`src/rfdetr/_namespace.py`	Forwards `oriented` into the legacy builder namespace.
`pyproject.toml`	Adds codespell ignore for “dota” and mypy override for new dataset module.
`tests/utilities/test_rotated_box_ops.py`	Unit tests for rotated ops and losses.
`tests/datasets/test_dota_detection.py`	Unit tests for DOTA parsing/dataset behavior.
`tests/models/test_obb_*`	Tests for oriented head, matcher/criterion, postprocess, and export.
`tests/training/test_obb_integration.py`	Integration tests for config/namespace wiring.

Copilot · 2026-04-10T19:36:30Z

src/rfdetr/datasets/dota_detection.py

+    if image_set == "train":
+        resize_wrappers = AlbumentationsWrapper.from_config(
+            [
+                {"Resize": {"height": resolution, "width": resolution}},
+            ]
+        )
+        aug_wrappers = AlbumentationsWrapper.from_config(
+            [
+                {"HorizontalFlip": {"p": 0.5}},
+                {"VerticalFlip": {"p": 0.5}},
+                {"RandomRotate90": {"p": 0.5}},
+            ]
+        )
+        return Compose([*resize_wrappers, *aug_wrappers, to_image, to_float, normalize])


make_dota_transforms() builds a pipeline out of AlbumentationsWrapper geometric transforms, but AlbumentationsWrapper only transforms targets when the target dict contains a "boxes" key (see rfdetr/datasets/transforms.py:642-645). DOTA targets only provide "corners" / "boxes_obb", so resize/flip/rotate will modify the image while leaving the geometry unchanged, producing incorrect boxes_obb after DotaNormalize (and potentially silently training on wrong labels). Consider either (a) extending AlbumentationsWrapper (or a DOTA-specific wrapper) to transform the 4 corner keypoints via Albumentations keypoint_params, or (b) avoiding Albumentations geometric transforms here and implementing corner-aware transforms (including consistent filtering when boxes become invalid).

Copilot · 2026-04-10T19:36:31Z

src/rfdetr/models/postprocess.py

+            results = [
+                {"scores": sc, "labels": lb, "boxes_obb": ob, "corners": cn}
+                for sc, lb, ob, cn in zip(scores, labels, obb, corners)
+            ]


When oriented=True, PostProcess.forward() returns early with boxes_obb/corners and skips the existing mask gathering/resizing logic. If segmentation_head is enabled alongside oriented boxes, postprocessing will drop pred_masks entirely. Consider handling the out_masks branch in oriented mode as well (gather the same top-K masks and resize to target_sizes), and include them in the returned dicts.

Suggested change

results = [

{"scores": sc, "labels": lb, "boxes_obb": ob, "corners": cn}

for sc, lb, ob, cn in zip(scores, labels, obb, corners)

]

results = []

if out_masks is not None:

for i in range(out_masks.shape[0]):

res_i = {

"scores": scores[i],

"labels": labels[i],

"boxes_obb": obb[i],

"corners": corners[i],

}

k_idx = topk_boxes[i]

masks_i = torch.gather(

out_masks[i],

0,

k_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, out_masks.shape[-2], out_masks.shape[-1]),

) # [K, Hm, Wm]

h, w = target_sizes[i].tolist()

masks_i = F.interpolate(

masks_i.unsqueeze(1),

size=(int(h), int(w)),

mode="bilinear",

align_corners=False,

) # [K,1,H,W]

res_i["masks"] = masks_i > 0.0

results.append(res_i)

else:

results = [

{"scores": sc, "labels": lb, "boxes_obb": ob, "corners": cn}

for sc, lb, ob, cn in zip(scores, labels, obb, corners)

]

Copilot · 2026-04-10T19:36:31Z

src/rfdetr/models/criterion.py

+        if self.oriented:
+            loss_bbox = F.l1_loss(src_boxes[..., :4], target_boxes[..., :4], reduction="none")
+            losses["loss_bbox"] = loss_bbox.sum() / num_boxes
+            loss_kld = kld_loss(src_boxes, target_boxes)
+            losses["loss_giou"] = loss_kld.sum() / num_boxes


In oriented mode, loss_boxes() stores the KLD regression term under the key "loss_giou". This keeps the existing weight_dict wiring working, but it makes logs/metrics misleading ("giou" is no longer GIoU). Consider additionally emitting a correctly named key like "loss_kld" (or similar) for logging/monitoring while retaining "loss_giou" for backward compatibility, or add an inline comment clarifying the semantic change.

Copilot · 2026-04-10T19:36:32Z

tests/datasets/test_dota_detection.py

+class TestBuildDota:
+    def test_builds_dataset(self, dota_root: Path) -> None:
+        import types
+
+        args = types.SimpleNamespace(dataset_dir=str(dota_root.parent))
+        root_with_split = dota_root.parent / "train"
+        root_with_split.mkdir(exist_ok=True)
+        (root_with_split / "images").mkdir(exist_ok=True)
+        (root_with_split / "labelTxt").mkdir(exist_ok=True)
+        img = Image.new("RGB", (50, 50), color="blue")
+        img.save(root_with_split / "images" / "test.png")
+        (root_with_split / "labelTxt" / "test.txt").write_text("")
+        args.dataset_dir = str(dota_root.parent)
+        dataset = build_dota("train", args, 256)
+        assert isinstance(dataset, DotaDetection)


TestBuildDota.test_builds_dataset only asserts that build_dota() returns a DotaDetection instance, but it never calls dataset[0]. Since build_dota() wires up a transform pipeline, this test currently won’t catch transform/target-sync issues (e.g., geometric transforms not updating corners / boxes_obb). Consider adding a __getitem__ assertion that verifies shapes and that boxes_obb stays in normalized [0,1] coords after transforms.

farukalamai added 10 commits April 3, 2026 04:37

feat: add rotated box utilities for OBB support

e394646

feat: add DOTA v1.0 dataset loader for OBB support

a70c7fb

feat: add oriented detection head with angle prediction MLP

caf9902

feat: add GWD matching cost and KLD loss for oriented boxes

07e4b9b

feat: add oriented box postprocessing with corner output

96f51f5

feat: verify oriented box ONNX export produces 5D output

9a3c6f6

feat: wire oriented flag through training pipeline and namespace

e7ed4ef

fix: rename dota.py to dota_detection.py and fix edge cases

d3b281f

fix: avoid in-place normalization and add edge case tests

ba227fe

fix: add oriented to BuilderArgs protocol, zero-init angle_embed, fix…

66d9413

… return type

farukalamai requested review from Borda, SkalskiP, isaacrob and probicheaux as code owners April 3, 2026 19:40

farukalamai added 5 commits April 4, 2026 01:44

fix: add dota to codespell ignore list and fix mypy Dataset subclass …

6ae1cce

…error

fix: resolve mypy type-arg error and codespell dota false positive

3381899

fix: provide generic type argument to Dataset instead of type-ignore

2d1d5d8

fix: add dota_detection to mypy ignore list matching other dataset mo…

9308416

…dules

Merge remote-tracking branch 'upstream/develop' into feat/obb-support

e81c9f4

Borda requested a review from Copilot April 10, 2026 19:30

Copilot started reviewing on behalf of Borda April 10, 2026 19:30 View session

Copilot AI reviewed Apr 10, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Oriented Bounding Box (OBB) support for rotated object detection#921

Add Oriented Bounding Box (OBB) support for rotated object detection#921
farukalamai wants to merge 15 commits intoroboflow:developfrom
farukalamai:feat/obb-support

farukalamai commented Apr 3, 2026

Uh oh!

codecov bot commented Apr 3, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

-            results = [
-                {"scores": sc, "labels": lb, "boxes_obb": ob, "corners": cn}
-                for sc, lb, ob, cn in zip(scores, labels, obb, corners)
-            ]
+            results = []
+            if out_masks is not None:
+                for i in range(out_masks.shape[0]):
+                    res_i = {
+                        "scores": scores[i],
+                        "labels": labels[i],
+                        "boxes_obb": obb[i],
+                        "corners": corners[i],
+                    }
+                    k_idx = topk_boxes[i]
+                    masks_i = torch.gather(
+                        out_masks[i],
+,
+                        k_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, out_masks.shape[-2], out_masks.shape[-1]),
+                    )  # [K, Hm, Wm]
+                    h, w = target_sizes[i].tolist()
+                    masks_i = F.interpolate(
+                        masks_i.unsqueeze(1),
+                        size=(int(h), int(w)),
+                        mode="bilinear",
+                        align_corners=False,
+                    )  # [K,1,H,W]
+                    res_i["masks"] = masks_i > 0.0
+                    results.append(res_i)
+            else:
+                results = [
+                    {"scores": sc, "labels": lb, "boxes_obb": ob, "corners": cn}
+                    for sc, lb, ob, cn in zip(scores, labels, obb, corners)
+                ]

Conversation

farukalamai commented Apr 3, 2026

What does this PR do?

Type of Change

Testing

Checklist

Additional Context

Files changed

Uh oh!

codecov bot commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov bot commented Apr 3, 2026 •

edited

Loading